import sys
import os
import torch
import random
import numpy as np
from tqdm import tqdm
import gymnasium as gym
import matplotlib.pyplot as plt
import torch.nn.functional as F

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from QNet import QNet
from rl_utils import ReplayBuffer, moving_average


class DQN:
    """DQN 算法"""

    def __init__(
        self,
        state_dim,
        hidden_dim,
        action_dim,
        learning_rate,
        gamma,
        epsilon,
        target_update_freq,
        device=None,
    ):
        # 不断更新的训练 Q 网络
        self.QNet = QNet(state_dim, hidden_dim, action_dim).to(device)
        # 目标 Q 网络
        self.tar_QNet = QNet(state_dim, hidden_dim, action_dim).to(device)
        # Adam 优化器
        self.optimizer = torch.optim.Adam(self.QNet.parameters(), lr=learning_rate)

        self.state_dim = state_dim
        self.action_dim = action_dim

        # 折扣因子
        self.gamma = gamma
        # epsilon-greedy 策略中的 epsilon
        self.epsilon = epsilon

        # 目标网络更新频率
        self.target_update_freq = target_update_freq
        # 记录更新次数
        self.update_count = 0

        self.device = device

    def take_action(self, state):
        """根据 epsilon-greedy 策略选择动作"""
        if np.random.random() < self.epsilon:
            action = np.random.randint(self.action_dim)
        else:
            state = torch.tensor(np.array([state]), dtype=torch.float).to(self.device)
            action = self.QNet(state).argmax().item()
        return action

    def update(self, state, action, reward, next_state, done):
        """
        更新 Q 网络参数

        state: torch.Tensor, 形状为 (batch_size, state_dim)
        action: torch.Tensor, 形状为 (batch_size, 1)
        reward: torch.Tensor, 形状为 (batch_size, 1)
        next_state: torch.Tensor, 形状为 (batch_size, state_dim)
        done: torch.Tensor, 形状为 (batch_size, 1)
        """

        # 将输入转化为 Tensor，并确保其形状正确
        state = torch.tensor(state, dtype=torch.float).view(-1, self.state_dim).to(self.device)
        action = torch.tensor(action, dtype=torch.long).view(-1, 1).to(self.device)
        reward = torch.tensor(reward, dtype=torch.float).view(-1, 1).to(self.device)
        next_state = (
            torch.tensor(next_state, dtype=torch.float).view(-1, self.state_dim).to(self.device)
        )
        done = torch.tensor(done, dtype=torch.float).view(-1, 1).to(self.device)

        # 获得训练网络预测的 Q(s,a) 值
        qsa = self.QNet(state).gather(1, action)

        # 计算 TD 目标 r + \gamma * max_a' Q_target(s', a')
        max_tar_qsa = self.tar_QNet(next_state).max(1)[0].view(-1, 1)
        tar_qsa = reward + self.gamma * max_tar_qsa * (1 - done)

        # 计算 MSE 损失函数
        loss = torch.mean(F.mse_loss(qsa, tar_qsa))

        # 梯度下降
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        self.update_count += 1

        # 每 target_update_freq 次更新后，更新目标网络参数
        if self.update_count % self.target_update_freq == 0:
            self.tar_QNet.load_state_dict(self.QNet.state_dict())


if __name__ == "__main__":
    lr = 2e-3
    epochs = 10
    num_episodes = 500
    hidden_dim = 128
    gamma = 0.98
    epsilon = 0.01
    target_update = 10
    buffer_size = 10000
    minimal_size = 500
    batch_size = 64
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

    env_name = "CartPole-v0"
    env = gym.make(env_name)
    env.reset(seed=0)
    random.seed(0)
    np.random.seed(0)
    torch.manual_seed(0)
    replay_buffer = ReplayBuffer(buffer_size)
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n

    agent = DQN(state_dim, hidden_dim, action_dim, lr, gamma, epsilon, target_update, device)

    episodes_per_epoch = int(num_episodes / epochs)
    return_list = []
    for e in range(epochs):
        with tqdm(total=int(episodes_per_epoch), desc="Epoch %d" % e) as pbar:
            for i_episode in tqdm(range(episodes_per_epoch)):
                episode_return = 0

                state, _ = env.reset()
                done = False
                while not done:
                    action = agent.take_action(state)
                    next_state, reward, terminated, truncated, _ = env.step(action)
                    done = terminated or truncated  # 合并 terminated 和 truncated

                    replay_buffer.add(state, action, reward, next_state, done)
                    state = next_state
                    episode_return += reward

                    # 经验回放池中的样本数量大于 minimal_size 时，才开始更新 Q 网络参数
                    if replay_buffer.size() > minimal_size:
                        b_s, b_a, b_r, b_ns, b_d = replay_buffer.sample(batch_size)
                        agent.update(b_s, b_a, b_r, b_ns, b_d)

                return_list.append(episode_return)
                if (i_episode + 1) % 10 == 0:
                    pbar.set_postfix(
                        {
                            "episode": "%d" % (num_episodes / epochs * e + i_episode + 1),
                            "return": "%.3f" % np.mean(return_list[-10:]),
                        }
                    )
                pbar.update(1)

    episodes_list = list(range(len(return_list)))
    plt.plot(episodes_list, return_list)
    plt.xlabel("Episodes")
    plt.ylabel("Returns")
    plt.title("DQN on {}".format(env_name))
    plt.show()

    mv_return = moving_average(return_list, 9)
    plt.plot(episodes_list, mv_return)
    plt.xlabel("Episodes")
    plt.ylabel("Returns")
    plt.title("DQN on {}".format(env_name))
    plt.show()
